# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: api\handlers\status\worker_master.py
__author__ = "Attila Gerendi (Sanyi)"
import uuid, logging
from sqlalchemy import select as sql_select
from api.classes.status import *
from api.handlers.status import ScanStatusResourceHandler
from db import Connection
from db.tables.workers import WorkerRow, WorkersTable, WorkerAuthorization, WorkerStatus
from helpers.http_clients.scan_api import ScanApiClient
from helpers.http_clients.exceptions import *
from server.http_responses.http_errors import Http409, Http404
from server.http_responses.http_status_codes import Http204
from helpers.licensing import get_local_license

logger = logging.getLogger("api.handlers.status.worker_master")


class WorkerMaster(ScanStatusResourceHandler):
    def __init__(self, current_shard="master", **_):
        super().__init__()
        self.current_shard = current_shard

    def register_worker(self, body, **_):
        worker_info = ClsWorkerRegistrationInfo(**body)
        logger.debug("attempting to connect to worker at: %s", worker_info.endpoint)
        client = ScanApiClient(
            worker_info.endpoint,
            worker_info.scanning_app,
            request_timeout=2,
            ssl_verify=worker_info.ssl_verify,
        )
        try:
            client.check_connection()
        except HTTP404:
            raise Http409("Scanning Engine response error (404)")
        except HTTP500:
            raise Http409("Scanning Engine internal error (500)")
        except HTTPException as e:
            print(e.result)
            raise Http409("Scanning Engine response error", details=e.result)
        except RequestError:
            raise Http409("Could not connect to the Scanning Engine")

        with Connection(self.current_shard) as (db):
            q = (
                sql_select((WorkerRow.worker_id,))
                .where(WorkerRow.endpoint == worker_info.endpoint)
                .where(WorkerRow.scanning_app == worker_info.scanning_app)
            )
            wi = db.execute(q).scalar()
            if wi:
                raise Http409("Endpoint already exists")
            worker_id = str(uuid.uuid4())
            worker_info.authorization = WorkerAuthorization.PENDING
            worker_info.worker_id = worker_id
            q = WorkersTable.insert().values(
                worker_id=worker_id,
                scanning_app=worker_info.scanning_app,
                authorization=worker_info.authorization,
                endpoint=worker_info.endpoint,
                api_key=worker_info.api_key,
                ssl_verify=worker_info.ssl_verify,
                description=worker_info.description,
            )
            db.execute(q)
        return worker_info

    def get_worker_status(self, worker_id, **_):
        with Connection(self.current_shard) as (db):
            q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http409("Scanning Engine not found")
            worker = dict(worker)
        logger.debug("attempting to connect to worker at: %s", worker["endpoint"])
        client = ScanApiClient(
            worker["endpoint"],
            worker["scanning_app"],
            request_timeout=2,
            ssl_verify=worker["ssl_verify"],
        )
        try:
            client.check_connection()
        except HTTP401:
            raise Http409("Scanning Engine authentication error")
        except HTTP404:
            raise Http409("Scanning Engine response error (404)")
        except HTTP500:
            raise Http409("Scanning Engine internal error (500)")
        except HTTPException as e:
            print(e.result)
            raise Http409("Scanning Engine response error", details=e.result)
        except RequestError:
            raise Http409("Could not connect to the Scanning Engine")

        return ClsWorkerInfo(
            scanning_app=worker["scanning_app"],
            endpoint=worker["endpoint"],
            worker_id=worker["worker_id"],
            ssl_verify=worker["ssl_verify"],
            description=worker["ssl_verify"],
            authorization=worker["authorization"],
        )

    def unregister_worker(self, worker_id, **_):
        with Connection(self.current_shard) as (db):
            q = (
                sql_select((WorkerRow,))
                .where(WorkerRow.worker_id == worker_id)
                .where(
                    WorkerRow.authorization.in_(
                        (WorkerAuthorization.PENDING, WorkerAuthorization.AUTHORIZED)
                    )
                )
            )
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http409("Scanning Engine not found")
            q = (
                (
                    WorkersTable.update(
                        values=dict(
                            authorization=WorkerAuthorization.DETACHED,
                            status=WorkerStatus.DETACHED,
                        )
                    )
                )
                .where(WorkerRow.worker_id == worker_id)
                .where(
                    WorkerRow.authorization.in_(
                        (WorkerAuthorization.PENDING, WorkerAuthorization.AUTHORIZED)
                    )
                )
            )
            db.execute(q)
            return Http204()

    def update_worker_status(self, worker_id, body, **_):
        worker_status = ClsWorkerStatus(**body)
        with Connection(self.current_shard) as (db):
            q = (
                sql_select((WorkerRow,))
                .where(WorkerRow.worker_id == worker_id)
                .where(
                    WorkerRow.authorization.in_(
                        (WorkerAuthorization.PENDING, WorkerAuthorization.AUTHORIZED)
                    )
                )
            )
            worker = db.execute(q).fetchone()
            if not worker:
                raise Http409("Scanning Engine not found")
            if worker_status.status != "online":
                q = (
                    WorkersTable.update(values=dict(status=worker_status.status))
                ).where(WorkerRow.worker_id == worker_id)
                db.execute(q)
                return Http204()
            worker = dict(worker)
        client = ScanApiClient(
            worker["endpoint"],
            worker["scanning_app"],
            request_timeout=2,
            ssl_verify=worker["ssl_verify"],
        )
        if (
            not worker_status.license
            or worker_status.license != get_local_license()["license_key"]
        ):
            activation_data = None
            try:
                client.post_request(
                    "/worker/add_licence", data=dict(activation_data=activation_data)
                )
            except HTTPException as e:
                raise Http409("Scanning Engine response error", details=e.result)
            except RequestError:
                raise Http409("Could not connect to the Scanning Engine")

        else:
            try:
                client.get_worker_status()
            except HTTPException as e:
                raise Http409("Scanning Engine response error", details=e.result)
            except RequestError:
                raise Http409("Could not connect to the Scanning Engine")

            with Connection(self.current_shard) as (db):
                q = (
                    WorkersTable.update(values=dict(status=worker_status.status))
                ).where(WorkerRow.worker_id == worker_id)
                db.execute(q)
            return Http204()
